#include <sys/time.h> 
#include <iostream>

#include "randomc.h"
#include "stocc.h"

#include "ranker.h"

#ifndef MULTIFILE_PROJECT
// If compiled as a single file then include these cpp files, 
// If compiled as a project then compile and link in these cpp files.
   #include "mersenne.cpp"             // code for random number generator
   #include "stoc1.cpp"                // random library source code
   #include "userintf.cpp"             // define system specific user interface
#endif

using namespace std;

void Yperm(float *Ytime, float *Ydelta, const int N)
{
  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  for (int n = 1; n < N; n++){
    int idx = sto.IRandom(0, n); // uniform over [0,n]

    float time = Ytime[n];
    float delta = Ydelta[n];

    Ytime[n] = Ytime[idx];
    Ydelta[n] = Ydelta[idx];

    Ytime[idx] = time;
    Ydelta[idx] = delta;
  }
}

void Index0(const float *Ydelta, int *N0, const int N)
{
  int count = 0;
  for (int n = 0; n < N; n++){
    if (Ydelta[n] == 1.0){
      N0[count] = n;
      count++;
    }
  }
}

void Den(float *den, const float *Y, const int *N0, const int N, const int n0)
{ 
  for(int n = 0; n < n0; n++){
    float val = 0.0;
    int ind = N0[n];
    for (int j = 0; j < N; j++){
      if(Y[j] >= Y[ind])
	val++;
    }

    den[n] = val;
  }
}


void Tstat(const float *Ytime, const float *X, float *T, const int *N0, const float *Den, const int N, const int K, const int n0)
{  
  for (int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    // survival test statistic
    float stat = 0.0;
    float se = 0.0;
    float den = 0.0;
    for (int n = 0; n < n0; n++){
      int ind = N0[n];
      float nom = 0.0;
      float nom2 = 0.0;
      for (int j = 0; j < N; j++){
	if(Ytime[j] >= Ytime[ind]){
	  float aux = X[iter+j];
	  nom += aux;
	  nom2 += aux*aux;
	}
      }

      den = Den[n];
      stat += X[iter+ind] - nom/den;
      se += nom2/den - nom*nom/den/den;
    }

    T[k] = stat*stat/se;
  }
}

void TstatTmax(const float *Ytime, const float *X, float *T, const int *N0, const float *Den, const int N, const int K, const int n0, float *Tmax)
{ 
  float zeta = 0.0;
  for (int k = 0; k < K; k++){
    // compute start of X segment
    int iter = k*N;

    // survival test statistic
    float stat = 0.0;
    float se = 0.0;
    float den = 0.0;
    for (int n = 0; n < n0; n++){
      int ind = N0[n];
      float nom = 0.0;
      float nom2 = 0.0;
      for (int j = 0; j < N; j++){
	if(Ytime[j] >= Ytime[ind]){
	  float aux = X[iter+j];
	  nom += aux;
	  nom2 += aux*aux;
	}
      }

      den = Den[n];
      stat += X[iter+ind] - nom/den;
      se += nom2/den - nom*nom/den/den;
    }

    T[k] = stat*stat/se;

    if (zeta < fabs(T[k]))
      zeta = fabs(T[k]);
  }

  Tmax[0] = zeta;
}

void Pvalue(const float *T, const float *Ttilde, float *p, float *P, const float *Tmax, const int K)
{
  for(int k = 0; k < K; k++){
    if(fabs(Ttilde[k]) >= fabs(T[k]))
      p[k]++;
    if(Tmax[0] >= fabs(T[k]))
      P[k]++;
  }
}

long int seed = (long int)time(0); // global seed

int main(const int argc, const char *argv[])
{
  // argv[1]: N
  // argv[2]: K
  // argv[3]: B

  extern long int seed;
  seed++;
  StochasticLib1 sto(seed);      // make instance of random library

  int N = atoi(argv[1]);
  int K = atoi(argv[2]);
  int B = atoi(argv[3]);

  // print
  std::cout<<"(N="<<N<<",K="<<K<<",B="<<B<<")"<<"\n";

  // allocate memory  
  float *X;
  X = (float *)malloc(N*K*sizeof(float));
  float *Ytime;
  Ytime = (float *)malloc(N*sizeof(float));
  float *Ydelta;
  Ydelta = (float *)malloc(N*sizeof(float));
  float *p;
  p = (float *)malloc(K*sizeof(float));
  float *P;
  P = (float *)malloc(K*sizeof(float));
  float *T;
  T = (float *)malloc(K*sizeof(float));
  float *Ttilde;
  Ttilde = (float *)malloc(K*sizeof(float));
  float *Tmax;
  Tmax = (float *)malloc(sizeof(float));
  Tmax[0] = 0.0;

  // generate X, Ydelta, Ttime
  int count = 0;
  for (int n = 0; n < N; n++){
    Ytime[n] = fabs(sto.Normal(0, 1.0));
    Ydelta[n] = sto.Bernoulli(0.7);
    for (int k = 0; k < K; k++){
      X[count] = sto.Normal(0, 1.0);
      count++;
    }
  }

  // compute ranks of X;
  string method = "average";  // Can also be "min" or "max" or "default"
  for (int k = 0; k < K; k++){
    vector<double> a(N);
    for (int n = 0; n < N; n++)
      a[n] = X[k*N+n];

    vector<double> ranks;
    rank(a, ranks, method);

    for (int n = 0; n < N; n++)
      X[k*N+n] = ranks[n];
  }

  // start time profiling
  timeval t1, t2;
  double elapsed;
  gettimeofday(&t1, NULL);

  for (int k = 0; k < K; k++){
    p[k] = 0;
    P[k] = 0;
    T[k] = 0.0;
    Ttilde[k] = 0.0;
  }

  // compute n0
  int n0 = 0;
  for (int n = 0; n < N; n++){
    if (Ydelta[n] == 1.0)
      n0++;
  }

  float *den;
  den = (float *)malloc(n0*sizeof(float));

  int *N0;
  N0 = (int *)malloc(n0*sizeof(int));

  // compute N0
  Index0(Ydelta, N0, N);

  // compute T
  Den(den, Ytime, N0, N, n0);
  Tstat(Ytime, X, T, N0, den, N, K, n0);

  for (int b = 0; b < B; b++){
    // permute Y
    Yperm(Ytime, Ydelta, N);

    // compute N0
    Index0(Ydelta, N0, N);

    // compute test statistic
    Den(den, Ytime, N0, N, n0);
    TstatTmax(Ytime, X, Ttilde, N0, den, N, K, n0, Tmax);

    // compute p-values
    Pvalue(T, Ttilde, p, P, Tmax, K);
  }

  gettimeofday(&t2, NULL);
  elapsed = (t2.tv_sec - t1.tv_sec) * 1000.0;      // sec to ms
  elapsed += (t2.tv_usec - t1.tv_usec) / 1000.0;

  std::cout<<"elapsed time:"<<elapsed<<" milliseconds"<<"\n";

  // cleanup
  free(Ytime);
  free(Ydelta);
  free(X);
  free(p);
  free(P);
  free(T);
  free(Ttilde);
  free(Tmax);
  free(N0);
  free(den);

  return(0);
}
